"""
WebSocket连接测试脚本
====================

测试WebSocket实时数据推送功能
"""

import asyncio
import json
import logging
from datetime import datetime
from typing import List

try:
    import websockets
except ImportError:
    print("❌ 需要安装websockets库")
    print("   运行: pip install websockets")
    exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class WebSocketTestClient:
    """WebSocket测试客户端"""
    
    def __init__(self, url: str):
        self.url = url
        self.websocket = None
        self.received_messages = []
    
    async def connect(self):
        """连接到WebSocket服务器"""
        try:
            self.websocket = await websockets.connect(self.url)
            logger.info(f"✅ 已连接到 {self.url}")
            return True
        except Exception as e:
            logger.error(f"❌ 连接失败: {e}")
            return False
    
    async def send_message(self, message: dict):
        """发送消息"""
        try:
            if not self.websocket:
                logger.error("❌ WebSocket未连接")
                return False
            
            message_str = json.dumps(message)
            await self.websocket.send(message_str)
            logger.info(f"📤 已发送消息: {message}")
            return True
        except Exception as e:
            logger.error(f"❌ 发送消息失败: {e}")
            return False
    
    async def receive_message(self, timeout: float = 5.0):
        """接收消息"""
        try:
            if not self.websocket:
                logger.error("❌ WebSocket未连接")
                return None
            
            message_str = await asyncio.wait_for(
                self.websocket.recv(),
                timeout=timeout
            )
            message = json.loads(message_str)
            self.received_messages.append(message)
            logger.info(f"📥 已接收消息: {message}")
            return message
        except asyncio.TimeoutError:
            logger.warning("⏰ 接收消息超时")
            return None
        except Exception as e:
            logger.error(f"❌ 接收消息失败: {e}")
            return None
    
    async def subscribe(self, symbols: List[str]):
        """订阅股票代码"""
        message = {
            "type": "subscribe",
            "data": {
                "symbols": symbols
            }
        }
        return await self.send_message(message)
    
    async def unsubscribe(self, symbols: List[str]):
        """取消订阅"""
        message = {
            "type": "unsubscribe",
            "data": {
                "symbols": symbols
            }
        }
        return await self.send_message(message)
    
    async def listen(self, duration: int = 30):
        """监听消息"""
        logger.info(f"👂 开始监听消息（持续 {duration} 秒）...")
        
        start_time = asyncio.get_event_loop().time()
        message_count = 0
        
        while True:
            elapsed = asyncio.get_event_loop().time() - start_time
            if elapsed >= duration:
                logger.info(f"⏱️  监听时间到（{duration}秒），停止监听")
                break
            
            remaining = duration - elapsed
            message = await self.receive_message(timeout=min(remaining, 5.0))
            
            if message:
                message_count += 1
                message_type = message.get("type", "unknown")
                logger.info(f"📬 收到消息 #{message_count} (类型: {message_type})")
        
        logger.info(f"✅ 监听完成，共收到 {message_count} 条消息")
        return message_count
    
    async def close(self):
        """关闭连接"""
        if self.websocket:
            await self.websocket.close()
            logger.info("🔌 已断开连接")
    
    def get_received_messages(self):
        """获取所有接收到的消息"""
        return self.received_messages


async def test_basic_connection(url: str):
    """测试基本连接"""
    logger.info("=" * 60)
    logger.info("测试1: 基本连接测试")
    logger.info("=" * 60)
    
    client = WebSocketTestClient(url)
    
    try:
        # 连接
        connected = await client.connect()
        if not connected:
            logger.error("❌ 测试失败：无法建立连接")
            return False
        
        # 等待连接确认消息
        message = await client.receive_message(timeout=5.0)
        if message and message.get("type") == "connect":
            logger.info("✅ 收到连接确认消息")
        
        await asyncio.sleep(2)
        
        logger.info("✅ 基本连接测试通过")
        return True
        
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)
        return False
    finally:
        await client.close()


async def test_subscription(url: str):
    """测试订阅功能"""
    logger.info("=" * 60)
    logger.info("测试2: 订阅功能测试")
    logger.info("=" * 60)
    
    client = WebSocketTestClient(url)
    
    try:
        # 连接
        if not await client.connect():
            return False
        
        # 等待连接消息
        await client.receive_message(timeout=5.0)
        
        # 订阅股票代码
        symbols = ["000001.SH", "399001.SZ", "399006.SZ"]
        logger.info(f"📝 订阅股票代码: {symbols}")
        
        success = await client.subscribe(symbols)
        if not success:
            logger.error("❌ 订阅请求发送失败")
            return False
        
        # 接收订阅确认
        response = await client.receive_message(timeout=5.0)
        if response and response.get("type") == "subscribe":
            logger.info("✅ 收到订阅确认")
        
        # 监听数据消息
        logger.info("👂 监听数据推送（10秒）...")
        count = await client.listen(duration=10)
        
        if count > 0:
            logger.info(f"✅ 订阅功能测试通过（收到 {count} 条数据）")
            return True
        else:
            logger.warning("⚠️  未收到数据推送（可能是正常的，如果没有订阅者或非交易时间）")
            return True
        
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)
        return False
    finally:
        await client.close()


async def test_unsubscription(url: str):
    """测试取消订阅"""
    logger.info("=" * 60)
    logger.info("测试3: 取消订阅测试")
    logger.info("=" * 60)
    
    client = WebSocketTestClient(url)
    
    try:
        # 连接
        if not await client.connect():
            return False
        
        # 等待连接消息
        await client.receive_message(timeout=5.0)
        
        # 订阅
        symbols = ["000001.SH"]
        await client.subscribe(symbols)
        await client.receive_message(timeout=5.0)
        
        # 等待几秒
        await asyncio.sleep(3)
        
        # 取消订阅
        logger.info(f"🚫 取消订阅: {symbols}")
        await client.unsubscribe(symbols)
        
        # 接收取消订阅确认
        response = await client.receive_message(timeout=5.0)
        if response and response.get("type") == "unsubscribe":
            logger.info("✅ 收到取消订阅确认")
        
        logger.info("✅ 取消订阅测试通过")
        return True
        
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)
        return False
    finally:
        await client.close()


async def test_multiple_clients(url: str):
    """测试多客户端连接"""
    logger.info("=" * 60)
    logger.info("测试4: 多客户端连接测试")
    logger.info("=" * 60)
    
    clients = []
    
    try:
        # 创建多个客户端
        num_clients = 3
        for i in range(num_clients):
            client = WebSocketTestClient(url)
            if await client.connect():
                clients.append(client)
                logger.info(f"✅ 客户端 #{i+1} 已连接")
                # 等待连接消息
                await client.receive_message(timeout=5.0)
        
        if len(clients) != num_clients:
            logger.error(f"❌ 只有 {len(clients)}/{num_clients} 个客户端成功连接")
            return False
        
        # 每个客户端订阅不同的股票
        symbols_list = [
            ["000001.SH"],
            ["399001.SZ"],
            ["399006.SZ"]
        ]
        
        for i, (client, symbols) in enumerate(zip(clients, symbols_list)):
            await client.subscribe(symbols)
            logger.info(f"📝 客户端 #{i+1} 订阅: {symbols}")
            await client.receive_message(timeout=5.0)
        
        # 监听一段时间
        await asyncio.sleep(5)
        
        logger.info(f"✅ 多客户端连接测试通过（{num_clients} 个客户端）")
        return True
        
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)
        return False
    finally:
        # 关闭所有客户端
        for i, client in enumerate(clients):
            await client.close()
            logger.info(f"🔌 客户端 #{i+1} 已断开")


async def test_long_connection(url: str, duration: int = 30):
    """测试长连接"""
    logger.info("=" * 60)
    logger.info(f"测试5: 长连接测试（{duration}秒）")
    logger.info("=" * 60)
    
    client = WebSocketTestClient(url)
    
    try:
        # 连接
        if not await client.connect():
            return False
        
        # 等待连接消息
        await client.receive_message(timeout=5.0)
        
        # 订阅
        symbols = ["000001.SH", "399001.SZ"]
        await client.subscribe(symbols)
        await client.receive_message(timeout=5.0)
        
        # 长时间监听
        count = await client.listen(duration=duration)
        
        logger.info(f"✅ 长连接测试通过（持续 {duration}秒，收到 {count} 条消息）")
        return True
        
    except Exception as e:
        logger.error(f"❌ 测试失败: {e}", exc_info=True)
        return False
    finally:
        await client.close()


async def run_all_tests(url: str):
    """运行所有测试"""
    logger.info("🚀 开始WebSocket测试")
    logger.info(f"📍 测试地址: {url}")
    logger.info("=" * 60)
    
    tests = [
        ("基本连接", lambda: test_basic_connection(url)),
        ("订阅功能", lambda: test_subscription(url)),
        ("取消订阅", lambda: test_unsubscription(url)),
        ("多客户端", lambda: test_multiple_clients(url)),
        ("长连接", lambda: test_long_connection(url, duration=30)),
    ]
    
    results = []
    
    for test_name, test_func in tests:
        try:
            result = await test_func()
            results.append((test_name, result))
            await asyncio.sleep(2)  # 测试间隔
        except Exception as e:
            logger.error(f"❌ 测试 '{test_name}' 执行失败: {e}")
            results.append((test_name, False))
    
    # 打印测试结果
    logger.info("=" * 60)
    logger.info("📊 测试结果汇总")
    logger.info("=" * 60)
    
    passed = sum(1 for _, result in results if result)
    total = len(results)
    
    for test_name, result in results:
        status = "✅ 通过" if result else "❌ 失败"
        logger.info(f"{status} - {test_name}")
    
    logger.info("=" * 60)
    logger.info(f"总计: {passed}/{total} 个测试通过")
    
    if passed == total:
        logger.info("🎉 所有测试通过！")
    else:
        logger.warning(f"⚠️  有 {total - passed} 个测试失败")
    
    return passed == total


async def main():
    """主函数"""
    # WebSocket URL配置
    # 直接连接到market_data_service
    direct_url = "ws://localhost:8004/api/v1/market/ws/market-data"
    # 通过API Gateway
    gateway_url = "ws://localhost:8000/ws/market-data"
    
    print("\n" + "=" * 60)
    print("WebSocket连接测试")
    print("=" * 60)
    print("\n请选择测试模式：")
    print("1. 直接连接到 Market Data Service (端口 8004)")
    print("2. 通过 API Gateway (端口 8000)")
    print("3. 运行全部测试")
    
    choice = input("\n请输入选择 (1/2/3): ").strip()
    
    if choice == "1":
        logger.info("📍 测试模式: 直接连接")
        url = direct_url
        await run_all_tests(url)
    elif choice == "2":
        logger.info("📍 测试模式: API Gateway")
        url = gateway_url
        await run_all_tests(url)
    elif choice == "3":
        logger.info("📍 测试模式: 全部测试")
        logger.info("\n🔹 测试1: 直接连接到 Market Data Service")
        result1 = await run_all_tests(direct_url)
        
        await asyncio.sleep(3)
        
        logger.info("\n🔹 测试2: 通过 API Gateway")
        result2 = await run_all_tests(gateway_url)
        
        logger.info("\n" + "=" * 60)
        logger.info("📊 最终结果")
        logger.info("=" * 60)
        logger.info(f"直接连接: {'✅ 通过' if result1 else '❌ 失败'}")
        logger.info(f"API Gateway: {'✅ 通过' if result2 else '❌ 失败'}")
    else:
        logger.error("❌ 无效的选择")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        logger.info("\n👋 测试已取消")
    except Exception as e:
        logger.error(f"❌ 测试执行失败: {e}", exc_info=True)

